import streamlit as st
import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.datasets import make_moons
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 设置页面布局为宽屏模式
st.set_page_config(layout="wide")

st.title("全连接神经网络（Fully Connected Network, FCN/MLP）")

# ================= 基础介绍 =================
with st.form("fc_form1"):
    st.markdown("## 一、全连接网络基础")
    st.markdown("""
    ### 1. 定义
    全连接层（Fully Connected Layer）是神经网络中最常见的层，每个神经元与上一层的所有神经元相连。

    ### 2. 数学表达
    对于输入向量 $x \in \mathbb{R}^n$，全连接层输出 $y \in \mathbb{R}^m$：
    $$ y = Wx + b $$
    其中：
    - $W \in \mathbb{R}^{m \times n}$ 为权重矩阵
    - $b \in \mathbb{R}^m$ 为偏置向量

    当加上激活函数 $f$ 后：
    $$ y = f(Wx + b) $$

    ### 3. 特点
    - **优点**：结构简单，能拟合复杂函数；适合分类、回归任务
    - **缺点**：参数量大，计算开销大；不擅长处理图像/序列等结构化数据
    - **常见应用**：MLP、逻辑回归的扩展、神经网络最后的分类层
    """)

    col1, col2 = st.columns([10, 1])
    with col2:
        submit = st.form_submit_button("我已学习")
    if submit:
        st.success("恭喜你，已掌握全连接网络基础")
        st.balloons()

# ================= 代码演示 =================
with st.form("fc_form2"):
    st.markdown("## 二、代码演示：在二维数据上训练一个MLP")

    # 数据集生成
    X, y = make_moons(n_samples=500, noise=0.2, random_state=42)
    X = StandardScaler().fit_transform(X)
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)

    X_train_t = torch.tensor(X_train, dtype=torch.float32)
    y_train_t = torch.tensor(y_train, dtype=torch.long)
    X_test_t = torch.tensor(X_test, dtype=torch.float32)
    y_test_t = torch.tensor(y_test, dtype=torch.long)

    # 定义模型
    class FCNet(nn.Module):
        def __init__(self, input_dim, hidden_dim, output_dim):
            super().__init__()
            self.fc1 = nn.Linear(input_dim, hidden_dim)
            self.fc2 = nn.Linear(hidden_dim, output_dim)
        def forward(self, x):
            x = torch.relu(self.fc1(x))
            return self.fc2(x)

    model = FCNet(input_dim=2, hidden_dim=8, output_dim=2)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.01)

    # 训练循环
    epochs = 100
    losses = []
    for epoch in range(epochs):
        optimizer.zero_grad()
        outputs = model(X_train_t)
        loss = criterion(outputs, y_train_t)
        loss.backward()
        optimizer.step()
        losses.append(loss.item())

    # ========== 可视化 Loss 曲线 ==========
    fig1, ax1 = plt.subplots()
    ax1.plot(range(epochs), losses, label="Train Loss")
    ax1.set_xlabel("Epoch")
    ax1.set_ylabel("Loss")
    ax1.set_title("训练过程中Loss下降曲线")
    ax1.legend()
    st.pyplot(fig1)

    # ========== 可视化分类边界 ==========
    xx, yy = np.meshgrid(np.linspace(-2, 2, 200), np.linspace(-1.5, 1.5, 200))
    grid = torch.tensor(np.c_[xx.ravel(), yy.ravel()], dtype=torch.float32)
    with torch.no_grad():
        preds = torch.argmax(model(grid), axis=1).numpy()
    Z = preds.reshape(xx.shape)

    fig2, ax2 = plt.subplots()
    ax2.contourf(xx, yy, Z, alpha=0.3, cmap=plt.cm.coolwarm)
    ax2.scatter(X_test[:,0], X_test[:,1], c=y_test, cmap=plt.cm.coolwarm, edgecolors="k")
    ax2.set_title("测试集分类效果")
    st.pyplot(fig2)

    st.code("""
# 定义模型
model = FCNet(input_dim=2, hidden_dim=8, output_dim=2)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)

# 训练循环
for epoch in range(100):
    optimizer.zero_grad()
    outputs = model(X_train_t)
    loss = criterion(outputs, y_train_t)
    loss.backward()
    optimizer.step()
""", language="python")

    col1, col2 = st.columns([10, 1])
    with col2:
        submit = st.form_submit_button("我已学习")
    if submit:
        st.success("恭喜你，已掌握全连接网络代码实现与效果")
        st.balloons()
